#!/usr/bin/env python
# -*- encoding: UTF-8 -*-

'''
This module contains some common routines used by other samples.
'''

# Python 2/3 compatibility
from __future__ import print_function
import sys
PY3 = sys.version_info[0] == 3

if PY3:
    from functools import reduce

from scipy import weave
import numpy as np
import cv2
import math

# built-in modules
import os
import itertools as it
from contextlib import contextmanager

from sql import Cam

import datetime

def datetime2str(dt):
    return str(dt) # '2016-12-30 11:12:46.332000'

def str2datetime(string):
    return datetime.datetime.strptime(string, '%Y-%m-%d %H:%M:%S.%f')

# my:
def plot_point(img,p2d,color=(0,0,0)):
    x,y = np.array(p2d).flatten()
    x = int(x)
    y = int(y)
    h,w = img.shape[:2]
    if (x > 0 and x < w) and (y > 0 and y < h):
        cv2.circle(img, (x,y), 10, color,5)

# 骨骼细化
def _thinningIteration(im, iter):
    I, M = im, np.zeros(im.shape, np.uint8)
    expr = """
    for (int i = 1; i < NI[0]-1; i++) {
        for (int j = 1; j < NI[1]-1; j++) {
            int p2 = I2(i-1, j);
            int p3 = I2(i-1, j+1);
            int p4 = I2(i, j+1);
            int p5 = I2(i+1, j+1);
            int p6 = I2(i+1, j);
            int p7 = I2(i+1, j-1);
            int p8 = I2(i, j-1);
            int p9 = I2(i-1, j-1);
            int A  = (p2 == 0 && p3 == 1) + (p3 == 0 && p4 == 1) +
                     (p4 == 0 && p5 == 1) + (p5 == 0 && p6 == 1) +
                     (p6 == 0 && p7 == 1) + (p7 == 0 && p8 == 1) +
                     (p8 == 0 && p9 == 1) + (p9 == 0 && p2 == 1);
            int B  = p2 + p3 + p4 + p5 + p6 + p7 + p8 + p9;
            int m1 = iter == 0 ? (p2 * p4 * p6) : (p2 * p4 * p8);
            int m2 = iter == 0 ? (p4 * p6 * p8) : (p2 * p6 * p8);
            if (A == 1 && B >= 2 && B <= 6 && m1 == 0 && m2 == 0) {
                M2(i,j) = 1;
            }
        }
    }
    """

    weave.inline(expr, ["I", "iter", "M"])
    return (I & ~M)

def thinning(src):
    dst = src.copy() / 255
    prev = np.zeros(src.shape[:2], np.uint8)
    diff = None

    while True:
        dst = _thinningIteration(dst, 0)
        dst = _thinningIteration(dst, 1)
        diff = np.absolute(dst - prev)
        prev = dst.copy()
        if np.sum(diff) == 0:
            break

    return dst * 255

# 直线提取:
def get_line(src,mtx,dist,win_name=None):
    dst = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY) # 灰度图
    dst = cv2.pyrDown(dst) # 降阶一次
    dst = cv2.pyrDown(dst) # 降阶两次
    retval,dst = cv2.threshold(dst,10,255,cv2.THRESH_BINARY_INV) # 二值化
    dst = cv2.morphologyEx(dst, cv2.MORPH_OPEN, np.ones((3,3),np.uint8)) # 开运算
    dst = cv2.morphologyEx(dst, cv2.MORPH_CLOSE, np.ones((3,3),np.uint8)) # 闭运算
    if win_name is not None: cv2.imshow(win_name+"dst",dst)
    dst = thinning(dst)
    lines = cv2.HoughLinesP(dst, 1, math.pi/180.0, 20, np.array([]), 10, 10)
    if lines is None: return None
    #print(lines)
    a,b,c = lines.shape
    max_d = 0
    line = None
    pt = []
    for i in range(a):
        for j in range(b):
            pt1 = lines[i][j][0:2]
            pt2 = lines[i][j][2:4]
            d = np.linalg.norm(pt1 - pt2)
            if d > max_d:
                max_d = d
                pt = [pt1,pt2]
                pointA = np.array(pt1*2**2,np.float64).reshape((1,1,2))
                pointB = np.array(pt2*2**2,np.float64).reshape((1,1,2))
                pointA = cv2.undistortPoints(pointA,mtx,dist,np.array([]),mtx) # 畸变矫正
                pointB = cv2.undistortPoints(pointB,mtx,dist,np.array([]),mtx) # 畸变矫正
                pointA = np.hstack((pointA.reshape((2,)),1))
                pointB = np.hstack((pointB.reshape((2,)),1))
                line = np.cross(pointA,pointB) # 直线的齐次表示
    if win_name is not None:
        cdst = cv2.cvtColor(dst, cv2.COLOR_GRAY2BGR)
        cv2.line(cdst, (pt[0][0], pt[0][1]), (pt[1][0],pt[1][1]), (0, 0, 255), 3)
        cv2.imshow(win_name+"line",cdst)
    return np.matrix(line.reshape((3,1))) # 直线的齐次表示


def __tr2flat(tr,flat):
    '''
    将投影到相机坐标系下的平面齐次表示变换到另一个坐标系
    需要考虑, 该平面一定过相机的原点,
    实际上就是平面的齐次表示的最后一个值,不应该仍然是 0
    假设变换矩阵为 T_C^W 表示将坐标从 C 下变换到 W
    则 C 的原点在 W 中的表示为 T_C^W * [0,0,0,1]^T
    下述函数中 pi 为平面在相机坐标系 C 下的表示
    tr 为从相机坐标系 C 到世界坐标系 W 的变换矩阵, T_C^W
    :param tr:
    :param flat:
    :return:
    '''
    flat = tr*flat
    if flat[-1] == 0:
        C = tr[:,3] # C 的原点在 W 中的表示
        flat[-1] = -C.transpose()*flat # 确保 C^T*pi=0 即确保相机的原点在新平面pi上
        assert (C.transpose()*flat == 0).all()
    return flat


def get_line2(src1,cam1,tr1,src2,cam2,tr2):
    '''
    从相机图像重建三维直线
    :param src1: 图像1
    :param mtx1: 相机1内参矩阵
    :param dist1: 相机1畸变参数
    :param tr1: 从世界坐标系到相机1坐标系变换矩阵,T_c_o 从opencv的标定函数获得
    :param src2: 图像2
    :param mtx2:
    :param dist2:
    :param tr2:
    :return: L,Lstar 在世界坐标系中的直线的 Pulucke两点表示和两面表示
    '''
    assert isinstance(cam1,Cam)
    assert isinstance(cam2,Cam)
    line1 = get_line(src1,cam1.mtx,cam1.dist)
    line2 = get_line(src2,cam2.mtx,cam2.dist)
    if line1 is None or line2 is None: return None,None
    flat1 = cam1.get_P(src1).transpose()*line1.reshape((3,1))
    flat2 = cam2.get_P(src2).transpose()*line2.reshape((3,1))
    # 此时求解出来的平面为相机坐标系下的平面表达式，且过相机原点，
    # 即  Ax+By+Cz=0 ，但(x,y,z)的单位仍为相机坐标系的单位 cam.unit
    # 变换到以 mm 为单位的坐标系，即为
    # A*x*unit + B*y*unit + C*z*unit = 0
    # 显然，平面的齐次坐标并没有变化，
    # 事实上，对于一个平面的齐次表达式，空间单位的变化仅影响最后一项
    # A*x + B*y + C*z + D = 0
    # A*x*unit + B*y*unit + C*z*unit + D = 0
    # A*x + B*y + C*z + D/unit = 0
    # 又因为在实际应用中，最后一项是由相机在世界坐标系下的坐标求得的
    # 所以，在求解平面时不必考虑单位缩放的问题，只要保证相机原点在世界坐标系下单位的正确性即可
    flat1 = __tr2flat(np.linalg.inv(tr1),flat1)
    flat2 = __tr2flat(np.linalg.inv(tr2),flat2)
    Lstar = flat1*flat2.transpose() - flat2*flat1.transpose() # 空间直线的两面 Plucker 表示
    L = dual_plucker(Lstar) # 空间直线的两点 Plucker 表示
    return L, Lstar

# 求直线的 plucker 表示 L 的对偶表示
def dual_plucker(L):
    """
    一般 plucker 表示是用两空间点的齐次坐标表示, L_{4,4} = AB^T - BA^T
    但也可以用两空间的平面的齐次坐标表示, L^* = pi1*pi2^T - pi2*pi1^T
    两者互为对偶关系: L^* x L = 0, 关系如下:
    l_12 : l_13 : l_14 : l_23 : l_42 : l_34 ==
    == l^*_34 : l^*_42 : l^*_23 : l^*_14 : l^*_13 : l^*_12
    :param L: 空间直线的 plucker 表示矩阵
    :return: 相同空间直线的 plucker 表示矩阵的对偶表示矩阵
    """
    dual_L = L.copy()
    dual_L[0,1] = L[2,3]
    dual_L[0,2] = L[3,1]
    dual_L[0,3] = L[1,2]
    dual_L[1,2] = L[0,3]
    dual_L[3,1] = L[0,2]
    dual_L[2,3] = L[0,1]
    return dual_L - dual_L.transpose()

def plotLine(ax,p0,p1,c):
    # TODO 从点 p0 到 p1 的直线
    ax.plot3D([p0[0],p1[0]], [p0[1],p1[1]], [p0[2],p1[2]], color=c)

def plotTf(ax,Tf,p0):
    # TODO 绘制Tf的原点坐标
    array = np.array
    matrix = np.matrix
    p0 = array(p0).flatten()
    p1 = Tf[:3,3]
    p1 = array(p1).flatten()
    plotLine(ax,p0,p1,'k') # 原点间连线为黑色
    # TODO 绘制Tf的坐标轴,x,y,z 分别为红绿蓝
    k = 20
    x = Tf * matrix([k,0,0,1]).transpose()
    y = Tf * matrix([0,k,0,1]).transpose()
    z = Tf * matrix([0,0,k,1]).transpose()
    plotLine(ax,p1,x,'r')
    plotLine(ax,p1,y,'g')
    plotLine(ax,p1,z,'b')
    return p1

#  显示棋盘格
def drawChessboard(src,w,h):
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    gray = cv2.cvtColor(src,cv2.COLOR_BGR2GRAY)
    ret, corners = cv2.findChessboardCorners(gray, (w,h),None)
    if ret:
        cv2.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
        cv2.drawChessboardCorners(src, (w,h), corners,ret)

# opencv: common.py
image_extensions = ['.bmp', '.jpg', '.jpeg', '.png', '.tif', '.tiff', '.pbm', '.pgm', '.ppm']

class Bunch(object):
    def __init__(self, **kw):
        self.__dict__.update(kw)
    def __str__(self):
        return str(self.__dict__)

def splitfn(fn):
    path, fn = os.path.split(fn)
    name, ext = os.path.splitext(fn)
    return path, name, ext

def anorm2(a):
    return (a*a).sum(-1)
def anorm(a):
    return np.sqrt( anorm2(a) )

def homotrans(H, x, y):
    xs = H[0, 0]*x + H[0, 1]*y + H[0, 2]
    ys = H[1, 0]*x + H[1, 1]*y + H[1, 2]
    s  = H[2, 0]*x + H[2, 1]*y + H[2, 2]
    return xs/s, ys/s

def to_rect(a):
    a = np.ravel(a)
    if len(a) == 2:
        a = (0, 0, a[0], a[1])
    return np.array(a, np.float64).reshape(2, 2)

def rect2rect_mtx(src, dst):
    src, dst = to_rect(src), to_rect(dst)
    cx, cy = (dst[1] - dst[0]) / (src[1] - src[0])
    tx, ty = dst[0] - src[0] * (cx, cy)
    M = np.float64([[ cx,  0, tx],
                    [  0, cy, ty],
                    [  0,  0,  1]])
    return M


def lookat(eye, target, up = (0, 0, 1)):
    fwd = np.asarray(target, np.float64) - eye
    fwd /= anorm(fwd)
    right = np.cross(fwd, up)
    right /= anorm(right)
    down = np.cross(fwd, right)
    R = np.float64([right, down, fwd])
    tvec = -np.dot(R, eye)
    return R, tvec

def mtx2rvec(R):
    w, u, vt = cv2.SVDecomp(R - np.eye(3))
    p = vt[0] + u[:,0]*w[0]    # same as np.dot(R, vt[0])
    c = np.dot(vt[0], p)
    s = np.dot(vt[1], p)
    axis = np.cross(vt[0], vt[1])
    return axis * np.arctan2(s, c)

def draw_str(dst, target, s):
    x, y = target
    cv2.putText(dst, s, (x+1, y+1), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 0), thickness = 2)
    cv2.putText(dst, s, (x, y), cv2.FONT_HERSHEY_PLAIN, 1.0, (255, 255, 255))

class Sketcher:
    def __init__(self, windowname, dests, colors_func):
        self.prev_pt = None
        self.windowname = windowname
        self.dests = dests
        self.colors_func = colors_func
        self.dirty = False
        self.show()
        cv2.setMouseCallback(self.windowname, self.on_mouse)

    def show(self):
        cv2.imshow(self.windowname, self.dests[0])

    def on_mouse(self, event, x, y, flags, param):
        pt = (x, y)
        if event == cv2.EVENT_LBUTTONDOWN:
            self.prev_pt = pt
        elif event == cv2.EVENT_LBUTTONUP:
            self.prev_pt = None

        if self.prev_pt and flags & cv2.EVENT_FLAG_LBUTTON:
            for dst, color in zip(self.dests, self.colors_func()):
                cv2.line(dst, self.prev_pt, pt, color, 5)
            self.dirty = True
            self.prev_pt = pt
            self.show()


# palette data from matplotlib/_cm.py
_jet_data =   {'red':   ((0., 0, 0), (0.35, 0, 0), (0.66, 1, 1), (0.89,1, 1),
                         (1, 0.5, 0.5)),
               'green': ((0., 0, 0), (0.125,0, 0), (0.375,1, 1), (0.64,1, 1),
                         (0.91,0,0), (1, 0, 0)),
               'blue':  ((0., 0.5, 0.5), (0.11, 1, 1), (0.34, 1, 1), (0.65,0, 0),
                         (1, 0, 0))}

cmap_data = { 'jet' : _jet_data }

def make_cmap(name, n=256):
    data = cmap_data[name]
    xs = np.linspace(0.0, 1.0, n)
    channels = []
    eps = 1e-6
    for ch_name in ['blue', 'green', 'red']:
        ch_data = data[ch_name]
        xp, yp = [], []
        for x, y1, y2 in ch_data:
            xp += [x, x+eps]
            yp += [y1, y2]
        ch = np.interp(xs, xp, yp)
        channels.append(ch)
    return np.uint8(np.array(channels).T*255)

def nothing(*arg, **kw):
    pass

def clock():
    return cv2.getTickCount() / cv2.getTickFrequency()

@contextmanager
def Timer(msg):
    print(msg, '...',)
    start = clock()
    try:
        yield
    finally:
        print("%.2f ms" % ((clock()-start)*1000))

class StatValue:
    def __init__(self, smooth_coef = 0.5):
        self.value = None
        self.smooth_coef = smooth_coef
    def update(self, v):
        if self.value is None:
            self.value = v
        else:
            c = self.smooth_coef
            self.value = c * self.value + (1.0-c) * v

class RectSelector:
    def __init__(self, win, callback):
        self.win = win
        self.callback = callback
        cv2.setMouseCallback(win, self.onmouse)
        self.drag_start = None
        self.drag_rect = None
    def onmouse(self, event, x, y, flags, param):
        x, y = np.int16([x, y]) # BUG
        if event == cv2.EVENT_LBUTTONDOWN:
            self.drag_start = (x, y)
            return
        if self.drag_start:
            if flags & cv2.EVENT_FLAG_LBUTTON:
                xo, yo = self.drag_start
                x0, y0 = np.minimum([xo, yo], [x, y])
                x1, y1 = np.maximum([xo, yo], [x, y])
                self.drag_rect = None
                if x1-x0 > 0 and y1-y0 > 0:
                    self.drag_rect = (x0, y0, x1, y1)
            else:
                rect = self.drag_rect
                self.drag_start = None
                self.drag_rect = None
                if rect:
                    self.callback(rect)
    def draw(self, vis):
        if not self.drag_rect:
            return False
        x0, y0, x1, y1 = self.drag_rect
        cv2.rectangle(vis, (x0, y0), (x1, y1), (0, 255, 0), 2)
        return True
    @property
    def dragging(self):
        return self.drag_rect is not None


def grouper(n, iterable, fillvalue=None):
    '''grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx'''
    args = [iter(iterable)] * n
    if PY3:
        output = it.zip_longest(fillvalue=fillvalue, *args)
    else:
        output = it.izip_longest(fillvalue=fillvalue, *args)
    return output

def mosaic(w, imgs):
    '''Make a grid from images.

    w    -- number of grid columns
    imgs -- images (must have same size and format)
    '''
    imgs = iter(imgs)
    if PY3:
        img0 = next(imgs)
    else:
        img0 = imgs.next()
    pad = np.zeros_like(img0)
    imgs = it.chain([img0], imgs)
    rows = grouper(w, imgs, pad)
    return np.vstack(map(np.hstack, rows))

def getsize(img):
    h, w = img.shape[:2]
    return w, h

def mdot(*args):
    return reduce(np.dot, args)

def draw_keypoints(vis, keypoints, color = (0, 255, 255)):
    for kp in keypoints:
            x, y = kp.pt
            cv2.circle(vis, (int(x), int(y)), 2, color)
